# Collect_Trajectory.py
import os
import cv2
import mss
import time
import h5py
import torch
import random
import argparse
import threading
import numpy as np
import pydirectinput
from collections import deque
from models.policy_net import MultiHeadPolicyNet
from Activate_Window import activate_window
from config import (FPS, DURATION, KNIGHT_HP_THRESHOLD, REWARD_RANGE, CHECK_TIME,WRITE_IN_FPS, DEVICE,
                    CAPTURE_WINDOW, KNIGHT_HP_WINDOW, BOSS_HP_WINDOW,BATTLE_WINDOW,
                    TRAJECTORY_SAVE_PATH, NUM_EPISODES, MODEL_PATH,
                    ATOMIC_ACTIONS, ACTION_DIMS)

# 全局动作状态（由主线程写入，动作线程读取）
current_action = {
    'ad'    : 0, # 0: none, 1: a, 2: d
    'space' : 0, # 0: up, 1: down
    'attack': 0, # 0: no attack, 1: j, 2: wjw, 3: sjs
    'dash'  : 0, # 0: no dash, 1: trigger dash
}
action_lock = threading.RLock() # 保护 current_action 的读写

# 全局退出事件，用于立即终止动作线程
shutdown_event = threading.Event()

# 全局小骑士血量参照图
# KNIGHT_REF_IMG_LAPTOP = cv2.imread("knight_ref_hp_laptop.png")
KNIGHT_REF_IMG = cv2.imread("../knight_ref_hp.png")

def ad_thread():
    """ 0-停止 1-左 2-右 """
    last_ad = 0
    while not shutdown_event.is_set():
        with action_lock:
            ad = current_action['ad']
        if ad != last_ad:
            if last_ad == 1:
                pydirectinput.keyUp('a')
            elif last_ad == 2:
                pydirectinput.keyUp('d')
            if ad == 1:
                pydirectinput.keyDown('a')
            elif ad == 2:
                pydirectinput.keyDown('d')
            last_ad = ad
        time.sleep(CHECK_TIME) # 每0.01秒，判断一次任务结束没有？没有结束的话，去current_action中确认最新指令并且执行

# todo:处理二段跳；对跳跃更细节的把控：落地的判断
def space_thread():
    """ 控制 space 跳跃 """
    last_space = 0
    while not shutdown_event.is_set():
        with action_lock:
            space = current_action['space']
        if space != last_space:
            if space == 1:
                pydirectinput.keyDown('space')
            else:
                pydirectinput.keyUp('space')
            last_space = space
        time.sleep(CHECK_TIME)

# todo:给执行下劈附加条件，不满足就执行正面劈砍
def attack_thread():
    """ 瞬时攻击动作（j / wjw / sjs）一直执行"""
    while not shutdown_event.is_set():
        attack_type = 0
        with action_lock:
            attack_type = current_action['attack'] # 只读取，不消费.实现连续攻击
        if attack_type == 1:
            pydirectinput.press('j')
        elif attack_type == 2:
            pydirectinput.keyDown('w')
            pydirectinput.press('j')
            pydirectinput.keyUp('w')
        elif attack_type == 3:
            pydirectinput.keyDown('s')
            pydirectinput.press('j')
            pydirectinput.keyUp('s')
        time.sleep(CHECK_TIME)

# todo:黑冲和普通冲；冲刺过程中对其他键的并行屏蔽；先按其他键再按冲刺的并行可行。
def dash_thread():
    while not shutdown_event.is_set():
        dash = 0
        with action_lock:
            dash = current_action['dash']
            if dash == 1:
                current_action['dash'] = 0  # 消费
        if dash == 1:
            # 暂时用冲刺后置来避免优先执行冲刺从而导致其他按键执行被打断
            time.sleep(WRITE_IN_FPS/(2*FPS))  # 75ms 延迟，让其他输入先生效
            pydirectinput.press('k')
        time.sleep(CHECK_TIME)

# 原始画面截图
def capture_frame(sct):
    img = np.array(sct.grab(CAPTURE_WINDOW))    # BGRA (H, W, 4)
    return img[:, :, :3]                        # BGR  (H, W, 3)

# 从原图提取小骑士、BOSS血量图、战斗图
def extract_subframe(frame,window):
    y  = window["top"]
    x  = window["left"]
    dy = window["height"]
    dx = window["width"]
    return frame[y:y+dy, x:x+dx]

# 小骑士血量图 -> 扣血数
def detect_knight_hp_change(cmp_img,threshold=KNIGHT_HP_THRESHOLD):
    match  = (cmp_img == KNIGHT_REF_IMG).all(axis=2)[0]
    diff   = np.diff(match.astype(np.int8), prepend=0, append=0)
    starts = np.where(diff == 1)[0]
    ends   = np.where(diff == -1)[0]
    change = np.sum(ends - starts >= threshold)
    return change

# BOSS血量图 -> 剩余血量百分比
def detect_boss_hp_percentage(hp_bar_img):
    # [5,0,96]
    if not np.array_equal(hp_bar_img[0, 0], [5, 0, 96]):
        return 0.0
    w = BOSS_HP_WINDOW["width"]
    diff = np.diff(hp_bar_img[0, :, :].astype(np.int16), axis=0)
    candidates = np.where(np.abs(diff).max(axis=1) > 1)[0]
    if len(candidates) > 0:
        return candidates[0] / w
    else:
        return 1.0

# 数据缓存器
class TrajectoryBuffer:
    def __init__(self, max_len):
        self.max_len = max_len
        self.counts     = deque(maxlen=max_len)
        self.frames     = deque(maxlen=max_len)
        self.rewards    = deque(maxlen=max_len)
        self.knight_hp  = deque(maxlen=max_len)
        self.boss_hp    = deque(maxlen=max_len)
        self.actions    = deque(maxlen=max_len)

    def push(self, count, frame, knight, boss, action):
        """ 推入新帧 """
        self.counts.append(int(count))
        self.frames.append(frame.copy())
        self.rewards.append(0)
        self.knight_hp.append(knight)
        self.boss_hp.append(boss)
        self.actions.append(np.array(action))

    def backfill_previous(self, new_reward):
        """ 将队列中所有帧的 reward 叠加 new_reward """
        start_idx = -1
        end_idx   = -min(len(self.rewards),REWARD_RANGE)
        for index in range(start_idx, end_idx - 1, -1):  # 从 -1 到 -4
            self.rewards[index] += new_reward

    def pop_oldest_and_save(self, h5_group):
        """ 弹出最老的一帧并保存 """
        if len(self.frames) == 0:
            return False

        # popleft() 弹出最左边（最老的）
        count   = self.counts.popleft()
        frame   = self.frames.popleft()
        reward  = self.rewards.popleft()
        knight  = self.knight_hp.popleft()
        boss    = self.boss_hp.popleft()
        action  = self.actions.popleft()

        # 写入 HDF5
        step_grp = h5_group.create_group(f"step_{count}")
        step_grp.create_dataset("count",     data=count)
        step_grp.create_dataset("frame",     data=frame) # compression="gzip"非必须,时间换空间
        step_grp.create_dataset("reward",    data=reward)
        step_grp.create_dataset("knight_hp", data=knight)
        step_grp.create_dataset("boss_hp",   data=boss)
        step_grp.create_dataset("action",    data=action)

        return True

    def is_full(self):
        return len(self.counts) == self.max_len

    def flush_all(self, h5_group):
        """ 游戏结束时，把剩下的全写进去 """
        while len(self.counts) > 0:
            self.pop_oldest_and_save(h5_group)

# 随机策略
def sample_random_action():
    action_ad       = random.choice(ACTION_DIMS['ad'])
    action_space    = random.choice(ACTION_DIMS['space'])
    action_dash   = random.choice(ACTION_DIMS['dash'])
    action_attack = random.choice(ACTION_DIMS['attack'])
    return [action_ad, action_space, action_dash, action_attack]

def release_all_keys():
    """
    释放所有可能被按下的游戏控制键
    防止 episode 结束后角色继续移动/跳跃等
    """
    for key in ATOMIC_ACTIONS:
        pydirectinput.keyUp(key)
    print("🎮 所有按键已释放")

def main(episode_id, save_path, model, device):
    complete_mark = True
    shutdown_event.clear()

    threads = [
        threading.Thread(target=ad_thread,      daemon=True),
        threading.Thread(target=space_thread,   daemon=True),
        threading.Thread(target=attack_thread,  daemon=True),
        threading.Thread(target=dash_thread,    daemon=True),
    ]
    for t in threads:
        t.start()

    activate_window("Hollow Knight")
    time.sleep(0.5)
    pydirectinput.press("w")
    time.sleep(2)
    pydirectinput.press("w")
    time.sleep(1)
    pydirectinput.press("space")
    time.sleep(3)

    with mss.mss() as sct, h5py.File(save_path, 'a') as f:
        group_name = f"episode_{episode_id}"
        if group_name in f:
            del f[group_name]
        episode_grp = f.create_group(group_name)

        buffer = TrajectoryBuffer(max_len=REWARD_RANGE+1)
        count            = 0             # 步数唯一标识
        prev_knight      = 10            # 小骑士上轮生命值
        prev_knight_loss = 0             # 小骑士上轮受伤值
        prev_boss        = 0             # BOSS上轮生命值
        recorded_action  = [0, 0, 0, 0]  # 小骑士上轮动作
        frame_stack      = deque(maxlen=WRITE_IN_FPS)

        start_time = time.time()
        while True:
            next_frame_time = start_time + (count + 1) / FPS
            current_time = time.time()
            if current_time < next_frame_time:
                time.sleep(next_frame_time - current_time)

            frame = capture_frame(sct)
            count += 1

            # 分割frame为三个部分
            # ①战斗画面
            # 预处理：截取，灰度，压缩
            battle_frame = extract_subframe(frame,BATTLE_WINDOW)
            gray_battle_frame = cv2.cvtColor(battle_frame, cv2.COLOR_BGR2GRAY)
            h = BATTLE_WINDOW["height"]
            w = BATTLE_WINDOW["width"]
            resized_gray_battle_frame = cv2.resize(gray_battle_frame, (w // 10, h // 10), interpolation=cv2.INTER_AREA)
            # todo:这里(w // 10, h // 10)可以考虑用固定大小例如(256,100)
            frame_stack.append(resized_gray_battle_frame)
            # 将灰度图特征提取,包含 位置、动作形态 信息,
            # 特征信息 合并上 按键信息，作为状态state传入策略网络，处理输出得到动作action
            # ②小骑士血量
            knight_hp_bar = extract_subframe(frame, KNIGHT_HP_WINDOW)
            current_knight_loss = detect_knight_hp_change(knight_hp_bar, KNIGHT_HP_THRESHOLD)
            current_knight = prev_knight - (current_knight_loss if current_knight_loss != prev_knight_loss else 0)
            if current_knight_loss != prev_knight_loss or prev_knight_loss == 0:
                prev_knight_loss = current_knight_loss
            # ③boss血量
            boss_hp_bar = extract_subframe(frame,BOSS_HP_WINDOW)
            current_boss = detect_boss_hp_percentage(boss_hp_bar)

            if count % WRITE_IN_FPS == 0:
                stacked_frames = np.stack(frame_stack)  # (WRITE_IN_FPS,100,256)
                # recorded_action = sample_random_action()
                recorded_action = model.get_action(stacked_frames,device=device)
                with action_lock:
                    current_action['ad']     = recorded_action[0]
                    current_action['space']  = recorded_action[1]
                    current_action['attack'] = recorded_action[2]
                    current_action['dash']   = recorded_action[3]
                buffer.push(count/WRITE_IN_FPS-1, stacked_frames, current_knight, current_boss, recorded_action)

            # 更新reward
            is_done = False
            # ①分析小骑士，其血量提供的对局信息更加确定
            if current_knight == prev_knight: # 无伤
                pass
            elif current_knight == 0:  # 角色死亡，采样结束，先不break等待boss血量结算
                buffer.backfill_previous(-5)
                print(f"step:{count:3d}:Knight-{prev_knight - current_knight}HP！dead！reward:-5")
                is_done = True
            elif current_knight < prev_knight: # 受伤惩罚比造成伤害更大，先学躲技能再学砍人
                buffer.backfill_previous(-5)
                print(f"step:{count:3d}:Knight-{prev_knight-current_knight}HP！reward:-5")
            # 这里没有考虑小骑士加血的情况，即在当前假设下血量是单调递减的
            else:
                is_done = True
                print("小骑士血量出错！")

            # ②分析BOSS血量，其血量提供的信息更迷惑多样
            if current_boss == prev_boss: # 开头还没攻击到、中间没攻击到
                pass
            elif current_boss == 0 and prev_boss != 0: # 单boss单阶段，这就是击败boss的瞬间
                if prev_boss > 0.1:
                    complete_mark = False  # 重新采样标记
                    print("血量采集出错，重新采集一次")
                else:
                    buffer.backfill_previous(1)
                    print(f"step:{count:3d}:BOSS-{prev_boss:.2f}HP BOSS dead！ reward:1")
                is_done = True
            elif (current_boss < prev_boss) or (prev_boss == 0 and current_boss != 0):
                buffer.backfill_previous(1)
                print(f"step:{count:3d}:BOSS-{(1 if prev_boss == 0 else prev_boss)-current_boss:.2f}HP reward:1")
            # 这里后面可以尝试处理的逻辑是：boss血量恢复了，部分回复或者回复到1.现在就只打单boss单阶段是没问题的
            else:
                is_done = True
                print("BOSS血量出错！")

            # reward更新完，如果 buffer 满了，就弹出最老的一帧保存到文件
            if buffer.is_full():
                buffer.pop_oldest_and_save(episode_grp)

            if is_done == True or (time.time() - start_time) >= DURATION:
                print(f"{time.time() - start_time:.1f}s : ⏹️ 采样结束")
                break

            prev_knight = current_knight
            prev_boss = current_boss

        total_time = time.time() - start_time
        actual_fps = count / total_time  # 帧数 / 总时间
        print(f"\n📊 采样统计:")
        print(f"   目标时长: {DURATION:.1f}s")
        print(f"   实际耗时: {total_time:.2f}s")
        print(f"   总帧数: {count}")
        print(f"   实际平均采样率: {actual_fps:.2f} FPS")

        if complete_mark == False:
            del f[group_name]  # 删除整个 episode
        else:
            # 把剩下的帧写进文件
            buffer.flush_all(episode_grp)

        release_all_keys()

    # 通知动作线程退出（立即生效）
    shutdown_event.set()  # 标记为退出状态

    for t in threads:
        t.join(timeout=0.5)

    return complete_mark

def get_next_episode_id(save_path=TRAJECTORY_SAVE_PATH):
    """ 读取 HDF5 文件中已有的 episode 数量，返回下一个 episode_id """
    try:
        with h5py.File(save_path, 'r') as f:
            # 找出所有以 'episode_' 开头的 group
            episode_ids = []
            for key in f.keys():
                if key.startswith("episode_"):
                    try:
                        idx = int(key.split("_")[1])
                        episode_ids.append(idx)
                    except ValueError:
                        continue  # 忽略无法解析的
            if len(episode_ids) == 0:
                return 0
            return max(episode_ids) + 1
    except (OSError, FileNotFoundError):
        # 如果文件不存在，说明是第一次运行
        return 0

if __name__ == '__main__':
    START_ID = get_next_episode_id(TRAJECTORY_SAVE_PATH)
    print(f"✅ 检测到当前最后一个 episode_id = {START_ID - 1}，将从 {START_ID} 开始采集")
    model = MultiHeadPolicyNet(history_len=WRITE_IN_FPS, action_dims=(3, 2, 4, 2))
    if os.path.exists(MODEL_PATH):
        model.load_state_dict(torch.load(MODEL_PATH))
    else:
        print("⚠️ 未找到模型，使用随机策略初始化")
    model.to(DEVICE)

    for i in range(NUM_EPISODES):
        episode_id = START_ID + i
        print(f"🎮 开始第 {episode_id + 1} 局采样 (episode_{episode_id})...")
        complete = False
        while complete == False:
            complete = main(episode_id=episode_id, save_path=TRAJECTORY_SAVE_PATH, model=model, device=DEVICE)
            time.sleep(15)  # 每局之间暂停